from itertools import count
from os import path

import torch

import DQN
from Agent import Agent
from env import FlappyEnvironment
from  matplotlib import pyplot as plt
import numpy as np

from utilenn import tensor_image_to_numpy_image

env = FlappyEnvironment()

model = DQN.DQN()

if torch.cuda.is_available():
    model.cuda()

if path.exists('./dqn.net'):
    model.load_state_dict(torch.load('./dqn.net'))

agent = Agent(model, 2)

env.reset()

env.step(1)


def showImage():
    # return tensor_image_to_numpy_image(env.current_frame)
    # return tensor_image_to_numpy_image(env.current_state)
    return env.raw_screen()


img = plt.imshow(showImage(), cmap='gray')

bestStep = 0

for c in range(5):
    env.reset()
    step = 0

    while True:
        action = agent.select_action(
            env.current_state,
            -1
        )
        done = env.step(action)

        img.set_data(showImage())

        plt.draw()
        plt.pause(0.0333)

        step += 1
        if done:
            break

    bestStep = max(bestStep, step)

print('best step', bestStep)
